package main

import (
	"context"
	"encoding/json"
	"log"
	"os"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/chromedp/chromedp"
)

type TrainInfo struct {
	DepartureTime string `json:"departure_time"`
	StartStation  string `json:"start_station"`
	ArrivalTime   string `json:"arrival_time"`
	EndStation    string `json:"end_station"`
	Duration      string `json:"duration"`
}

func main() {
	// 创建上下文
	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", true),
		chromedp.UserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"),
	)
	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
	defer cancel()

	ctx, cancel := chromedp.NewContext(allocCtx)
	defer cancel()

	// 设置超时时间
	ctx, cancel = context.WithTimeout(ctx, 20*time.Second)
	defer cancel()

	url := "https://trains.ctrip.com/webapp/train/list?ticketType=0&dStation=%E4%B8%8A%E6%B5%B7&aStation=%E5%8C%97%E4%BA%AC&dDate=2025-04-17&rDate=&hubCityName=&highSpeedOnly=0"

	log.Println("正在访问页面:", url)
	var htmlContent string
	err := chromedp.Run(ctx,
		chromedp.Navigate(url),
		chromedp.WaitReady(`body`, chromedp.ByQuery),
		chromedp.Sleep(2*time.Second),
		chromedp.ActionFunc(func(ctx context.Context) error {
			log.Println("开始查找#main元素...")
			return nil
		}),
		chromedp.WaitVisible(`#main`, chromedp.ByQuery),
		chromedp.ActionFunc(func(ctx context.Context) error {
			log.Println("已找到#main元素，开始获取内容")
			return nil
		}),
		chromedp.OuterHTML(`#main`, &htmlContent),
	)
	if err != nil {
		log.Fatalf("页面加载失败: %v", err)
	}

	log.Println("页面加载成功，开始解析数据...")
	// 使用goquery解析HTML
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(htmlContent))
	if err != nil {
		log.Fatal(err)
	}

	var trains []TrainInfo

	// 修改查找逻辑，直接查找.list-item
	doc.Find(`.train-content section[role="product"] .list-item`).Each(func(i int, s *goquery.Selection) {
	    var train TrainInfo
	
	    // 提取出发时间
	    train.DepartureTime = strings.TrimSpace(s.Find(".list-bd .from .time").Text())
	    
	    // 提取始发站
	    train.StartStation = strings.TrimSpace(s.Find(".list-bd .from .station").Text())
	    
	    // 提取到达时间
	    train.ArrivalTime = strings.TrimSpace(s.Find(".list-bd .to .time").Text())
	    
	    // 提取终点站
	    train.EndStation = strings.TrimSpace(s.Find(".list-bd .to .station").Text())
	    
	    // 提取运行时长并转换为hm格式
	    duration := strings.TrimSpace(s.Find(".list-bd .mid .haoshi").Text())
	    train.Duration = convertDuration(duration)
	
	    trains = append(trains, train)
	})

	// 将数据写入JSON文件
	file, err := os.Create("train.json")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	encoder.SetIndent("", "  ")
	err = encoder.Encode(trains)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("数据已成功写入train.json")
}

// 将"6时45分"转换为"6h45m"格式
func convertDuration(duration string) string {
	duration = strings.Replace(duration, "时", "h", 1)
	duration = strings.Replace(duration, "分", "m", 1)
	return duration
}
